#include <linux/init.h>  //_initdata
#include <linux/kallsyms.h>
#include <linux/lsm_hooks.h>
#include <linux/module.h>
#include <linux/printk.h>
#include <linux/types.h>  // umode_t

#define MAX_RO_PAGES 1024
static struct page* ro_pages[MAX_RO_PAGES];
static unsigned int ro_pages_len;

struct security_hook_heads probe_dummy_security_hook_heads;

#define MY_HOOK_INIT(HEAD, HOOK)                                               \
    {                                                                          \
        .head = &probe_dummy_security_hook_heads.HEAD, .hook = {.HEAD = HOOK } \
    }

// int (*inode_mkdir)(struct inode *dir, struct dentry *dentry, umode_t mode);
int linx_inode_mkdir(struct inode* dir, struct dentry* dentry, umode_t mode)
{
    printk("run linx_inode_mkdir:%s\n", dentry->d_name.name);
    return 0;
}

// int (*inode_rmdir)(struct inode *dir, struct dentry *dentry);
int linx_inode_rmdir(struct inode* dir, struct dentry* dentry)
{
    printk("run linx_inode_rmdir:%s\n", dentry->d_name.name);
    return 0;
}
// int (*inode_unlink)(struct inode *dir, struct dentry *dentry);
int linx_inode_unlink(struct inode* dir, struct dentry* dentry)
{
    printk("run linx_inode_unlink:%s\n", dentry->d_name.name);
    return 0;
}

static struct security_hook_list linx_hooks[] = {
    MY_HOOK_INIT(inode_mkdir, linx_inode_mkdir),
    MY_HOOK_INIT(inode_rmdir, linx_inode_rmdir),
    MY_HOOK_INIT(inode_unlink, linx_inode_unlink),
};

static int lsm_test_page_ro(void* addr)
{
    unsigned int i;
    int          unused;
    struct page* page;
    page = (struct page*)lookup_address((unsigned long)addr, &unused);
    if (!page) return 1;
    if (test_bit(_PAGE_BIT_RW, &(page->flags))) return 0;

    for (i = 0; i < ro_pages_len; i++)
        if (page == ro_pages[i]) return 0;
    if (ro_pages_len == MAX_RO_PAGES) return 1;
    ro_pages[ro_pages_len++] = page;
    return 0;
}

static int check_ro_pages(struct security_hook_heads* hooks)
{
    int i;
    for (i = 0; i < ARRAY_SIZE(linx_hooks); i++) {
        struct hlist_head* head = linx_hooks[i].head;
        if (lsm_test_page_ro(&head->first)) return 1;
    }
    return 0;
}

static inline void add_hook(struct security_hook_list* hook)
{
    hlist_add_tail_rcu(&hook->list, hook->head);
}

static inline void del_hook(struct security_hook_list* hook)
{
    hlist_del_rcu(&hook->list);
}

static void add_hooks(void)
{
    int idx;
    for (idx = 0; idx < ro_pages_len; idx++)
        set_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));

    for (idx = 0; idx < ARRAY_SIZE(linx_hooks); idx++)
        add_hook(&linx_hooks[idx]);

    for (idx = 0; idx < ro_pages_len; idx++)
        clear_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));
    return;
}

static void del_hooks(void)
{
    int idx;
    for (idx = 0; idx < ro_pages_len; idx++)
        set_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));

    for (idx = 0; idx < ARRAY_SIZE(linx_hooks); idx++)
        del_hook(&linx_hooks[idx]);

    for (idx = 0; idx < ro_pages_len; idx++)
        clear_bit(_PAGE_BIT_RW, &(ro_pages[idx]->flags));
    return;
}

static int proc_init(void)
{
    int                         idx;
    struct security_hook_heads* hooks =
        (struct security_hook_heads*)kallsyms_lookup_name(
            "security_hook_heads");
    if (!hooks) {
        printk(KERN_EMERG "hooks search failed\n");
        return -EINVAL;
    }
    for (idx = 0; idx < ARRAY_SIZE(linx_hooks); idx++)
        linx_hooks[idx].head =
            ((void*)hooks) + ((unsigned long)linx_hooks[idx].head) -
            ((unsigned long)&probe_dummy_security_hook_heads);

    if (check_ro_pages(hooks)) {
        printk(KERN_EMERG
               "Can't update security_hook_heads due to write protected.\n");
        return -EINVAL;
    }

    add_hooks();

    printk("hook-test install\n");
    return 0;
}

static void proc_cleanup(void)
{
    del_hooks();
    printk("hook-test uninstall.\n");
}

module_init(proc_init);
module_exit(proc_cleanup);
MODULE_LICENSE("GPL");
